#include <ntifs.h>

#include "Tool/Base.h"

EXTERN_C PVOID obHandle;

EXTERN_C HANDLE gHandle;

VOID Driver_Unlaod(PDRIVER_OBJECT pdriver)
{
	
	//UnRegisterCommucationCallBack();
	//UnProtectProcess();
	//UnInsertUserTimer();
	
}




NTSTATUS NTAPI DriverDispatch(PCommPackage package) {
	PVOID data = (PVOID)package->inData;
	NTSTATUS status = STATUS_UNSUCCESSFUL;
	switch (package->cmd)
	{
	case CMD_TEST: {
		status = STATUS_SUCCESS;
		break;
	}
	case CMD_GET_MOUDULE: {
		PModuleInfo pModuleInfo = (PModuleInfo)data;
		if (pModuleInfo) {
			SIZE_T imageSize = 0;
			ULONG_PTR imageBase = 0;
			status = GetModuleInfo((HANDLE)pModuleInfo->pid, (CHAR*)pModuleInfo->moduleName, (PULONG_PTR)&imageBase, &imageSize);
			if (NT_SUCCESS(status)) {
				pModuleInfo->moduleBase = imageBase;
				pModuleInfo->moduleSize = imageSize;
				status = STATUS_SUCCESS;
			}

		}
		break;
	}
	case CMD_READ_MEMORY: {
		PREADWRITEINFO info = (PREADWRITEINFO)data;
		if (info) {
			status = ReadMmeory((HANDLE)info->pid, (PVOID)info->targetAddress, (PVOID)info->buffer, info->size, READWRITECR3);
		}
		break;
	}
	case CMD_WRITE_MEMORY: {
		PREADWRITEINFO info = (PREADWRITEINFO)data;
		if (info) {
			status = WriteMemory((HANDLE)info->pid, (PVOID)info->targetAddress, (PVOID)info->buffer, info->size);
		}
		break;
	}
	case CMD_QUERY_MEMORY: {
	
		PQUERYINFO info = (PQUERYINFO)data;
		if (info) {
			status = QueryMemory((HANDLE)info->pid, (PVOID)info->targetAddress, (PMEMORY_BASIC_INFORMATION)info->memoryInfo);
		}
		break;
	}
	case CMD_PROCESS_PROCESS: {
		PPROTECTEDPROCESSINFO info = (PPROTECTEDPROCESSINFO)data;
		if (info) {
			status = ProtectProcessByPid((HANDLE)info->pid);
		}
		break;
	}
	case CMD_REMOTE_CALL: {

		PREMOTECALLINFO info = (PREMOTECALLINFO)data;
		if (info) {
			status =RemoteCall((HANDLE)info->pid,(PVOID)info->shellCode,info->shellCodeSize);
		}
		break;
	}
	default:
		status = STATUS_NOT_IMPLEMENTED;
		break;
	}
	package->retStatus = status;
	return status;
}


VOID CleanDriverEnv(PDRIVER_OBJECT pDriver, PUNICODE_STRING pReg) {
	NTSTATUS status = DriverCleanRegistry(pReg);
	if (!NT_SUCCESS(status)) {
		LOG("自清除注册表失败\r\n");
	}
	PKLDR_DATA_TABLE_ENTRY ldr = (PKLDR_DATA_TABLE_ENTRY)pDriver->DriverSection;
	if (ldr) {
		status = DriverSelfDelete(ldr->FullDllName.Buffer);
		if (!NT_SUCCESS(status)) {
			LOG("自删除驱动失败\r\n");
		}
	}
}


typedef NTSTATUS
(NTAPI
*pfnNtOpenProcess)(
	_Out_ PHANDLE ProcessHandle,
	_In_ ACCESS_MASK DesiredAccess,
	_In_ POBJECT_ATTRIBUTES ObjectAttributes,
	_In_opt_ PCLIENT_ID ClientId
);


pfnNtOpenProcess oldNtOpenProcess = NULL;

NTSTATUS
NTAPI
MyNtOpenProcess(
	_Out_ PHANDLE ProcessHandle,
	_In_ ACCESS_MASK DesiredAccess,
	_In_ POBJECT_ATTRIBUTES ObjectAttributes,
	_In_opt_ PCLIENT_ID ClientId
) {

	
	return oldNtOpenProcess(ProcessHandle, DesiredAccess, ObjectAttributes, ClientId);
}




PVOID GetFuncAddress(WCHAR* name) {
	UNICODE_STRING funName;
	RtlInitUnicodeString(&funName, name);
	return MmGetSystemRoutineAddress(&funName);
}


ULONG_PTR GetFunc() {

	return 0;
}

VOID FindComunicationFunc() {

	char pattern[] = "48******44******8B*****FF05****FF15****488BD8";;
	PUCHAR funAddress = (PUCHAR)searchCode("fltMgr.sys", ".text", pattern, 0);
	long offset = *(LONG*)(funAddress + 0x3);
	ULONG_PTR targetAddress = (ULONG_PTR)(funAddress + 0x7 + offset);
	LOG("funcAddress  = %llx ", targetAddress);
}






EXTERN_C NTSTATUS DriverEntry(PDRIVER_OBJECT pDriver, PUNICODE_STRING pReg)
{
	
	
	
	pDriver->DriverUnload = Driver_Unlaod;
	//ULONG_PTR address =  (ULONG_PTR)GetSssdtFcunAddress(GetUserGetForegrounWindowIndex());

	
	/*SetCreateThreadCallback(NULL);
	InjectDllX64((HANDLE)2596);*/
	//map view 读写
	//PVOID address = MapVirtualForPhyAddressByPid((HANDLE)2596, 0x00000000);

	//UnMapVirtualForPhyAddress((HANDLE)2596, 0x00000000);
	DbgBreakPoint();
	FindComunicationFunc();

	return STATUS_SUCCESS;
}